import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from einops import rearrange


class ObjectPointCloudLoss(nn.Module):
    def __init__(self):
        """
        Args:
            alpha: 重建损失权重
            beta: 时序约束权重
            dynamic_threshold: 是否使用动态阈值（基于真实运动量）
        """
        super().__init__()
        self.chamfer_weight = 1.0
        self.smooth_weight = 0.1
        self.dynamic_threshold = True
        
        # 静态阈值参数（需根据数据尺度调整）
        self.base_threshold = 0.1  # 基础阈值
        self.min_threshold = 0 # 最小运动量约束

    def chamfer_distance(self, x, y):
        """逐样本计算Chamfer距离"""
        dist = torch.cdist(x, y)  # (B, N, M)
        
        # x到y的最小距离
        min_dist_xy = torch.min(dist, dim=2)[0]  # (B, N)
        cost_xy = min_dist_xy.mean(dim=1)  # (B,)
        
        # y到x的最小距离
        min_dist_yx = torch.min(dist, dim=1)[0]  # (B, M)
        cost_yx = min_dist_yx.mean(dim=1)  # (B,)
        
        return (cost_xy + cost_yx)  # 保持样本维度 (B,)

    def get_smoothness_threshold(self, gt):
        """动态计算运动量阈值"""
        if self.dynamic_threshold:
            # 计算真实序列的相邻帧平均运动量
            gt_t = gt[:, :-1]  # (B, T-1, N, 3)
            gt_t1 = gt[:, 1:]  # (B, T-1, N, 3)
            real_movement = self.chamfer_distance(
                gt_t.reshape(-1, gt.shape[2], 3),
                gt_t1.reshape(-1, gt.shape[2], 3)
            ).mean()
            return torch.clamp(real_movement*0.5, min=self.min_threshold)
        else:
            return self.base_threshold

    def temporal_hinge_loss(self, pred, gt):
        """带阈值的时序平滑约束"""
        B, T, N, C = pred.shape
        
        # 获取相邻帧对
        pred_t = pred[:, :-1].reshape(-1, N, C)  # (B*(T-1), N, 3)
        pred_t1 = pred[:, 1:].reshape(-1, N, C)
        
        # 计算预测运动量
        pred_movement = self.chamfer_distance(pred_t, pred_t1)  # (B*(T-1),)
        
        # 计算动态阈值
        threshold = self.get_smoothness_threshold(gt) # 标量
        
        # Hinge Loss设计
        under_threshold = torch.clamp(threshold - pred_movement, min=0) # 运动小于 threshold，拉大运动
        over_threshold = torch.clamp(pred_movement - threshold*3, min=0) # 运动大于 threshold*3，减小运动
        
        # 组合约束：阻止过小运动同时限制过大突变
        return (under_threshold + over_threshold) / 2
    
    def temporal_chamfer_wo_smooth_loss(self, pred, gt):
        """带阈值的时序平滑约束"""
        B, T, N, C = pred.shape
        
        # 获取相邻帧对
        pred_t = pred[:, :-1].reshape(-1, N, C)  # (B*(T-1), N, 3)
        pred_t1 = pred[:, 1:].reshape(-1, N, C)
        # 计算预测运动量
        pred_movement = self.chamfer_distance(pred_t, pred_t1)  # (B*(T-1),)
        # 计算真值
        gt_t = gt[:, :-1].reshape(-1, N, C)  # (B*(T-1), N, 3)
        gt_t1 = gt[:, 1:].reshape(-1, N, C)
        # 计算预测运动量
        gt_movement = self.chamfer_distance(gt_t, gt_t1)  # (B*(T-1),)
        chamfer_wo_smooth_loss = F.l1_loss(pred_movement, gt_movement, reduction='none')
        chamfer_wo_smooth_loss = chamfer_wo_smooth_loss.contiguous().view(B, -1).mean(dim=-1)
        return chamfer_wo_smooth_loss

    def forward(self, pred, gt):
        # [b, t, v, c]
        # 逐帧重建损失
        B, T, N, C = pred.shape
        reconstruction_loss = self.chamfer_distance(
            pred.reshape(-1, N, C),
            gt.reshape(-1, N, C)
        ).contiguous().view(B, -1).mean(dim=-1)
        # 时序平滑约束
        # smoothness_loss = self.temporal_hinge_loss(pred, gt).mean()
        smoothness_loss = self.temporal_chamfer_wo_smooth_loss(pred, gt).contiguous().view(B, -1).mean(dim=-1)
        return {
            "loss": self.chamfer_weight * reconstruction_loss + self.smooth_weight * smoothness_loss,
            "loss_obj_chamfer": reconstruction_loss,
            "loss_obj_smooth": smoothness_loss,
        }


class HandJointLoss(nn.Module):
    def __init__(self):
        """
        Args:
            alpha: 重建损失权重
            beta: 时序约束权重
            dynamic_threshold: 是否使用动态阈值（基于真实运动量）
        """
        super().__init__()
        self.pos_weight = 1
        self.vel_weight = 0.2
        self.acc_weight = 0.05
    
    def forward(self, pred, gt):
        # [b, t, v, c]
        B, T, V, C = pred.shape
        pred_hand_traj = rearrange(pred, 'b t (m k) c -> b t m k c', m=2) # [b, t, 2, 21, 3]
        gt_hand_traj = rearrange(gt, 'b t (m k) c -> b t m k c', m=2)
        # temporal_weights = self.temporal_weights.to(device=pred_traj.device, dtype=pred_traj.dtype)
        pos_loss = F.smooth_l1_loss(pred_hand_traj, gt_hand_traj, reduction='none').mean(dim=[1, 2, 3, 4]) #  * temporal_weights[None, :, None, None]

        velocity = pred_hand_traj[:, 1:] - pred_hand_traj[:, :-1]
        gt_velocity = gt_hand_traj[:, 1:] - gt_hand_traj[:, :-1]
        vel_match_loss = F.l1_loss(velocity, gt_velocity, reduction='none').mean(dim=[1, 2, 3, 4])

        # acceleration = velocity[:, 1:] - velocity[:, :-1]
        # gt_acceleration = gt_velocity[:, 1:] - gt_velocity[:, :-1]
        # acc_match_loss = F.l1_loss(acceleration, gt_acceleration, reduction='none').mean(dim=[1, 2, 3, 4])

        return {
            "loss": self.pos_weight * pos_loss + self.vel_weight * vel_match_loss, #  + self.acc_weight * acc_match_loss,
            "loss_hand_pos": pos_loss,
            "loss_hand_vel": vel_match_loss,
            # "loss_hand_acc": acc_match_loss,
        }